package nsfs

import (
	"GLibHac/crypto"
	"GLibHac/fssystem"
	"GLibHac/tools"
	"encoding/hex"
	"fmt"
	"github.com/pkg/errors"
)

type XciPartitionType string

func (x XciPartitionType) String() string {
	switch x {
	case XciPartitionTypeUpdate:
		return "Update"
	case XciPartitionTypeNormal:
		return "Normal"
	case XciPartitionTypeSecure:
		return "Secure"
	case XciPartitionTypeLogo:
		return "Logo"
	case XciPartitionTypeRoot:
		return "Root"
	default:
		return fmt.Sprintf("XciPartitionType(%s)", string(x))
	}
}

const (
	XciPartitionTypeUpdate = XciPartitionType("update")
	XciPartitionTypeNormal = XciPartitionType("normal")
	XciPartitionTypeSecure = XciPartitionType("secure")
	XciPartitionTypeLogo   = XciPartitionType("logo")
	XciPartitionTypeRoot   = XciPartitionType("root")
)

type GameCardAttribute byte

func (t GameCardAttribute) String() string {
	switch t {
	case GameCardAttributeNone:
		return "None"
	case GameCardAttributeAutoBootFlag:
		return "AutoBootFlag"
	case GameCardAttributeHistoryEraseFlag:
		return "HistoryEraseFlag"
	case GameCardAttributeRepairToolFlag:
		return "RepairToolFlag"
	case GameCardAttributeDifferentRegionCupToTerraDeviceFlag:
		return "DifferentRegionCupToTerraDeviceFlag"
	case GameCardAttributeDifferentRegionCupToGlobalDeviceFlag:
		return "DifferentRegionCupToGlobalDeviceFlag"
	case GameCardAttributeHasCa10CertificateFlag:
		return "HasCa10CertificateFlag"
	default:
		return fmt.Sprintf("GameCardAttribute(%d)", t)
	}
}

const (
	GameCardAttributeNone                                 = GameCardAttribute(0)
	GameCardAttributeAutoBootFlag                         = GameCardAttribute(1)
	GameCardAttributeHistoryEraseFlag                     = GameCardAttribute(1 << 1)
	GameCardAttributeRepairToolFlag                       = GameCardAttribute(1 << 2)
	GameCardAttributeDifferentRegionCupToTerraDeviceFlag  = GameCardAttribute(1 << 3)
	GameCardAttributeDifferentRegionCupToGlobalDeviceFlag = GameCardAttribute(1 << 4)
	GameCardAttributeHasCa10CertificateFlag               = GameCardAttribute(1 << 7)
)

type GameCardSize byte

func (g GameCardSize) String() string {
	switch g {
	case GameCardSize1Gb:
		return "1GB"
	case GameCardSize2Gb:
		return "2GB"
	case GameCardSize4Gb:
		return "4GB"
	case GameCardSize8Gb:
		return "8GB"
	case GameCardSize16Gb:
		return "16GB"
	case GameCardSize32Gb:
		return "32GB"
	default:
		return fmt.Sprintf("GameCardSize(%d)", g)
	}
}

const (
	GameCardSize1Gb  = GameCardSize(0xFA)
	GameCardSize2Gb  = GameCardSize(0xF8)
	GameCardSize4Gb  = GameCardSize(0xF0)
	GameCardSize8Gb  = GameCardSize(0xE0)
	GameCardSize16Gb = GameCardSize(0xE1)
	GameCardSize32Gb = GameCardSize(0xE2)
)

const XciHeaderMagic = "HEAD"
const XciSignatureSize = 0x100
const GcTitleKeyKekIndexMax = 0x10

type XciHeader struct {
	baseStorage             fssystem.IStorage
	keySet                  *KeySet
	HasInitialData          bool
	magic                   string
	RomAreaStartPage        int32
	RackupAreaStartPage     int32
	KekIndex                byte
	TitleKeyDecIndex        byte
	GameCardSize            GameCardSize
	CardHeaderVersion       int32
	Flags                   GameCardAttribute
	PackageId               uint64
	ValidDataEndPage        uint64
	RootPartitionOffset     int64
	RootPartitionHeaderSize int64
	RootPartitionHeaderHash string
	SelSec                  int32
	SelT1Key                int32
	SelKey                  int32
	LimAreaPage             int32
	AesCbcIv                []byte
	InitialDataHash         []byte
	Signature               []byte
	ImageHash               []byte
	InitialDataAuthNonce    []byte
	InitialDataAuthMac      []byte
	InitialDataAuthData     []byte
	InitialDataPackageId    []byte
	DecryptedTitleKey       []byte
	InitialData             []byte
}

func NewXciHeader(keySet *KeySet, storage fssystem.IStorage) (this *XciHeader, err error) {
	this = &XciHeader{}
	this.baseStorage = storage
	this.keySet = keySet
	keyAreaStorage, bodyStorage, err := DetermineXciSubStorages(storage)
	if err != nil {
		return nil, err
	}
	if keyAreaStorage != nil {
		r := keyAreaStorage.AsStorage()
		this.HasInitialData = true
		this.InitialDataPackageId, err = r.ReadBytes(8)
		_, err = r.ReadBytes(8)
		if err != nil {
			return
		}
		this.InitialDataAuthData, err = r.ReadBytes(0x10)
		if err != nil {
			return
		}
		this.InitialDataAuthMac, err = r.ReadBytes(0x10)
		if err != nil {
			return
		}
		this.InitialDataAuthNonce, err = r.ReadBytes(0xC)
		if err != nil {
			return
		}
		err = r.SetPosition(0)
		if err != nil {
			return
		}
		this.InitialData, err = r.ReadBytes(512)
		if err != nil {
			return
		}
	}
	reader := bodyStorage.AsStorage()
	this.Signature, err = reader.ReadBytes(XciSignatureSize)
	if err != nil {
		return
	}
	tmpData, err := reader.ReadBytes(4)
	if err != nil {
		return nil, err
	}
	this.magic = string(tmpData)
	if XciHeaderMagic != this.magic {
		return nil, errors.New("invalid XCI file: Header magic invalid")
	}
	err = reader.SetPosition(XciSignatureSize)
	if err != nil {
		return
	}
	tmpData, err = reader.ReadBytes(XciSignatureSize)
	if err != nil {
		return
	}
	err = reader.SetPosition(XciSignatureSize + 4)
	if err != nil {
		return
	}
	this.RomAreaStartPage, err = reader.ReadInt()
	if err != nil {
		return
	}
	this.RackupAreaStartPage, err = reader.ReadInt()
	if err != nil {
		return
	}
	var tmpByte byte
	tmpByte, err = reader.ReadByte()
	if err != nil {
		return
	}
	this.KekIndex = tmpByte >> 4
	this.TitleKeyDecIndex = tmpByte & 7
	tmpByte, err = reader.ReadByte()
	if err != nil {
		return
	}
	this.GameCardSize = GameCardSize(tmpByte)
	tmpByte, err = reader.ReadByte()
	if err != nil {
		return
	}
	this.CardHeaderVersion = int32(tmpByte)
	tmpByte, err = reader.ReadByte()
	if err != nil {
		return
	}
	this.Flags = GameCardAttribute(tmpByte)
	this.PackageId, err = reader.ReadUint64()
	if err != nil {
		return
	}
	this.ValidDataEndPage, err = reader.ReadUint64()
	if err != nil {
		return
	}
	this.AesCbcIv, err = reader.ReadBytes(0x10)
	if err != nil {
		return
	}
	tools.ReverseSlice(this.AesCbcIv)
	this.RootPartitionOffset, err = reader.ReadInt64()
	if err != nil {
		return
	}
	this.RootPartitionHeaderSize, err = reader.ReadInt64()
	if err != nil {
		return
	}
	tmpData, err = reader.ReadBytes(0x20)
	if err != nil {
		return
	}
	this.RootPartitionHeaderHash = hex.EncodeToString(tmpData)
	this.InitialDataHash, err = reader.ReadBytes(0x20)
	if err != nil {
		return
	}
	this.SelSec, err = reader.ReadInt()
	if err != nil {
		return
	}
	this.SelT1Key, err = reader.ReadInt()
	if err != nil {
		return
	}
	this.SelKey, err = reader.ReadInt()
	if err != nil {
		return
	}
	this.LimAreaPage, err = reader.ReadInt()
	if err != nil {
		return
	}
	this.ImageHash = crypto.Sha256(tmpData)
	err = reader.SetPosition(this.RootPartitionOffset)
	if err != nil {
		return
	}
	tmpData, err = reader.ReadBytes(int(this.RootPartitionHeaderSize))
	if err != nil {
		return
	}
	actualHeaderHash := hex.EncodeToString(crypto.Sha256(tmpData))
	if actualHeaderHash != this.RootPartitionHeaderHash {
		return nil, errors.Errorf("invalid XCI, header hash=%s, actual=%s", this.RootPartitionHeaderHash, actualHeaderHash)
	}
	key := make([]byte, 0x10)
	if ok, err := DecryptCardInitialData(key, this.InitialData, this.KekIndex, keySet); err != nil {
		return nil, err
	} else if ok {
		this.DecryptedTitleKey = key
	}
	return
}

func DecryptCardInitialData(dest []byte, initialData []byte, kekIndex byte, keySet *KeySet) (bool, error) {
	if initialData == nil || len(initialData) == 0 {
		return false, nil
	}
	if kekIndex >= GcTitleKeyKekIndexMax {
		return false, nil
	}
	if keySet.GcTitleKeyKeks[kekIndex].IsZero() {
		return false, nil
	}
	return false, errors.New("not support decryption")
}

func DetermineXciSubStorages(baseStorage fssystem.IStorage) (fssystem.IStorage, fssystem.IStorage, error) {
	storageSize := baseStorage.Size()
	if storageSize >= 0x1104 {
		magic, err := baseStorage.ReadBytes(4)
		if err != nil {
			return nil, nil, err
		}
		if string(magic) == XciHeaderMagic {
			keyAreaStorage := fssystem.NewSubStorage(baseStorage, 0, 0x1000)
			bodyStorage := fssystem.NewSubStorage(baseStorage, 0x1000, storageSize-0x1000)
			return keyAreaStorage, bodyStorage, nil
		}
	}
	return nil, baseStorage.AsStorage(), nil
}

type Xci struct {
	*KeySet
	Header        *XciHeader
	baseStorage   fssystem.IStorage
	rootPartition *PartitionFileSystem
}

func (xci *Xci) GetRootPartition() *PartitionFileSystem {
	return xci.rootPartition
}

func (xci *Xci) OpenPartition(t XciPartitionType) (*PartitionFileSystem, error) {
	if t == XciPartitionTypeRoot {
		return xci.rootPartition, nil
	}
	file, err := xci.rootPartition.OpenFile(string(t), OpenFileModeRead)
	if err != nil {
		return nil, err
	}
	return NewPartitionFileSystem(file)
}

func NewXci(keySet *KeySet, reader fssystem.IStorage) (*Xci, error) {
	header, err := NewXciHeader(keySet, reader)
	if err != nil {
		return nil, err
	}
	if header.HasInitialData {
		reader = reader.Slice(0x1000)
	}
	partition, err := NewPartitionFileSystem(reader.Slice(header.RootPartitionOffset))
	if err != nil {
		return nil, err
	}
	return &Xci{
		KeySet:        keySet,
		Header:        header,
		baseStorage:   reader,
		rootPartition: partition,
	}, nil
}
